from fastapi import FastAPI, HTTPException, File, UploadFile
from pydantic import BaseModel
import os
from typing import Optional
import PyPDF2
import re
import uvicorn
import io

app = FastAPI()

class FilePath(BaseModel):
    path: str

@app.post("/pdftranchapter/process")
async def process(file: UploadFile = File(...)):
    """
    处理上传的PDF文件并返回生成文件目录地址
    """
    try:
        # 创建临时目录用于存储输出文件
        output_dir = "/home/huangjunren/workflow-fastapi/output"
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            
        # 读取上传的PDF文件
        content = await file.read()
        # 此处需要引入io库，否则会报错：name 'io' is not defined
        pdf_reader = PyPDF2.PdfReader(io.BytesIO(content))
            
        # 提取所有文本
        text = ''
        for page in pdf_reader.pages:
            text += page.extract_text()
        
        # 使用正则表达式查找章节
        chapter_pattern = r'\n第\s*\d+\s*章\s*'
        chapter_positions = [m.start() for m in re.finditer(chapter_pattern, text)]
        
        # 处理每个章节
        for i in range(len(chapter_positions)):
            # 获取章节标题
            next_pos = chapter_positions[i+1] if i < len(chapter_positions)-1 else len(text)
            chapter_content = text[chapter_positions[i]:next_pos]
            
            # 提取章节号
            chapter_num = re.search(r'第\s*(\d+)\s*章', chapter_content).group(1)
            
            # 创建输出文件
            output_file = os.path.join(output_dir, f'chapter_{chapter_num}.txt')
            
            # 写入文件
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(chapter_content.strip())
        
        return {
            "status": "success",
            "output_directory": output_dir
        }
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# 启动服务器
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=6000)